from Foundation import *
from AppKit import *

import csv
import itertools
import util

# from google.appengine.ext.webapp import template

def join(ctl, filename):
    out = csv.writer(open(filename, 'w'))
    for record in ctl.correlations():
        for x, row in zip(record['nodes'], record['matrix']):
            out.writerows(
                [record['state'], x, y, rho]
                for y, rho in zip(record['nodes'], row)
                if x < y)

def matrix(ctl, filename):
    out = csv.writer(open(filename, 'w'))
    for record in ctl.correlations():
        out.writerow([record['state']] + record['nodes'])
        out.writerows(
            [n] + row for n, row in 
            itertools.izip(record['nodes'], record['matrix']))

class Heatmap(NSView):
    def initWithFrame_(self, frame):
        self = super(Heatmap, self).initWithFrame_(frame)
        return self

    def drawRect_(self, rect):
        mps = NSMutableParagraphStyle.new()
        mps.setAlignment_(NSRightTextAlignment)
        right_align = { NSParagraphStyleAttributeName: mps }

        mps = NSMutableParagraphStyle.new()
        mps.setAlignment_(NSCenterTextAlignment)
        font = NSFont
        center_align = { NSParagraphStyleAttributeName: mps, 
                         NSUnderlineStyleAttributeName: NSSingleUnderlineStyle }

        rgb = NSColorSpace.genericRGBColorSpace()

        for state in self.params['states']:
            xf1 = NSAffineTransform.transform()
            xf1.translateXBy_yBy_(0, self.params['height'] - state['y'])
            xf1.concat()

            NSString.stringWithString_(state['name']).drawInRect_withAttributes_(
                NSMakeRect(0, -120, 240, 20), center_align)

            xf2 = NSAffineTransform.transform()
            xf2.translateXBy_yBy_(240, -240)
            xf2.concat()
            for node in state['nodes']:
                rc = NSMakeRect(-240, -node['offset'], 240, 20)
                NSString.stringWithString_(node['name']).drawInRect_withAttributes_(
                    rc, right_align)
            
            xf3 = NSAffineTransform.transform()
            xf3.translateXBy_yBy_(20, 20)
            xf3.rotateByDegrees_(90)
            xf3.concat()
            
            for node in state['nodes']:
                rc = NSMakeRect(0, -node['offset'], 240, 20)
                NSString.stringWithString_(node['name']).drawInRect_withAttributes_(
                    rc, {})

            xf3.invert()
            xf3.concat()
            
            xf3 = NSAffineTransform.transform()
            xf3.translateXBy_yBy_(-20, 20)
            xf3.scaleXBy_yBy_(20, 20)
            xf3.concat()
            
            for i, row in util.with_counter(state['matrix']):
                xf4 = NSAffineTransform.transform()
                xf4.translateXBy_yBy_(0, - i - 1)
                xf4.concat()
                for j, color in util.with_counter(row):
                    NSColor.colorWithCalibratedRed_green_blue_alpha_(
                        color['r'], color['g'], color['b'], 1.0).setFill()
                    NSRectFill(NSMakeRect(j + 1, 0, 1, 1))
                xf4.invert()
                xf4.concat()

            xf3.invert()
            xf3.concat()

            xf2.invert()
            xf2.concat()

            xf1.invert()
            xf1.concat()

def heatmap(ctl, filename):
    params = prepare_heatmap(ctl.correlations(), ctl.threshold)
    rect = NSMakeRect(0, 0, params['width'], params['height'])
    view = Heatmap.alloc().initWithFrame_(rect)
    view.params = params
    data = NSMutableData.dataWithCapacity_(100000)
    printOp = NSPrintOperation.PDFOperationWithView_insideRect_toData_(view, rect, data)
    printOp.runOperation()
    data.writeToFile_atomically_(filename, FALSE)

def prepare_heatmap(correl, threshold):
    params = {'width': 0, 'height': 0, 'states': []}
    for record in correl:
        side = 240 + 20 * len(record['nodes'])
        heatmap = [[color(x, threshold) for x in row] for row in record['matrix']]
        state = {'y': params['height'], 'name': record['state'], 'nodes': [],
                 'matrix': heatmap}
        params['width'] = max(params['width'], side)
        params['height'] += side
        for ix, n in util.with_counter(record['nodes']):
            state['nodes'].append({'name': n, 'offset': ix * 20})
        params['states'].append(state)
        params['width'] += 10
        params['height'] += 10
    return params

def heatmap_svg(correl, filename):
    open(filename, 'w').write(template.render('heatmap.svg', prepare_heatmap(correl)))

def color(x, threshold):
    if abs(x) < threshold:
        r, g, b = 1, 1, 1
    elif x >= 0:
        r, g, b = 1 - x, 1, 1 - x
    else:
        r, g, b = 1, 1 + x, 1 + x
    rr, gg, bb = (min(255, int(y * 256)) for y in (r, g, b))
    return { 'hex': '#%02x%02x%02x' % (rr, gg, bb),
             'r': r, 'g': g, 'b': b }
